package com.sinosoft.config.jdbc.dataSource;

import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.util.ReflectionUtils;

import javax.sql.DataSource;
import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Auther: zouren
 * @Date: 2019/5/10 15:29
 * @Description:
 */
@Slf4j
public class MyRoutingDataSource extends AbstractRoutingDataSource {
    private AtomicInteger count1 = new AtomicInteger();
    private AtomicInteger count2 = new AtomicInteger();
    @Override
    protected Object determineCurrentLookupKey() {

        String typeKey = DataSourceContextHolder.getJdbcType();
        Object re = null;
        // 写
        if (typeKey.equals(DataSourceType.write.getType())) {

            re =  getWriteKey();
        }else{
            re =  getReadKey();
        }
        log.info("chage:"+re.toString());
        return  re;

    }
    private Object getWriteKey(){
        if(writeListSize>=2){
            // 读简单负载均衡
            int number = count1.getAndAdd(1);
            if(count1.get()>9999){
                count1.set(-1);
            }
            int lookupKey = number % writeListSize;
            return  writeList.get(lookupKey);
        }
        return  writeList.get(1);

    }
    private Object getReadKey(){
        if(readListSize>=2){
            // 读简单负载均衡
            int number = count2.getAndAdd(1);
            if(count2.get()>=9999){
                count2.set(-1);
            }
            int lookupKey = number % readListSize;
            return  readList.get(lookupKey);
        }
        return  readList.get(1);
    }

    private  List<Object> writeList = new ArrayList<Object>();
    private  List<Object> readList = new ArrayList<Object>();
    private int writeListSize = 0;
    private int readListSize = 0;
    @Override
    public void afterPropertiesSet() {
        super.afterPropertiesSet();
        // 由于父类的resolvedDataSources属性是私有的子类获取不到，需要使用反射获取
        Field field = ReflectionUtils.findField(AbstractRoutingDataSource.class, "resolvedDataSources");
        // 设置可访问
        field.setAccessible(true);
        try {
            Map<Object, DataSource> resolvedDataSources = (Map<Object, DataSource>) field.get(this);
            resolvedDataSources.keySet().stream().forEach(o ->{
                if(o.toString().endsWith("_"+DataSourceType.write.getType())){
                    writeList.add(o);
                }else {
                    readList.add(o);
                }
            } );
            writeListSize = writeList.size();
            readListSize = readList.size();
        } catch (Exception e) {
            log.error("afterPropertiesSet error! ", e);
        }
    }
}
